#include <caml/alloc.h>
#include <caml/memory.h>
#include <caml/mlvalues.h>
#include <caml/callback.h>
#include "lbfgs.h"
#include <stdio.h>
#include <math.h>
#include <string.h>


void val_and_grad(value* f, double* ok_x, const double* x, double* g, 
        int n, double step, double* pfx)
{
    CAMLparam0();
    CAMLlocal3(vx, vg, vret);
    int x_okay = 1;
    int i;
    /* Copy x into a caml array. sizeof is to handle 64/32-bit architectures. */
    vx = alloc(n*sizeof(double)/sizeof(void*), Double_array_tag);
    for (i = 0; i < n; i++) { 
        if (!finite(x[i])) { x_okay = 0; }
        Store_double_field(vx, i, x[i]); 
    } 
    /* Save x if it's finite */
    if (x_okay) {
        for (i = 0; i < n; i++) {
            ok_x[i] = x[i];
        }
    }
    /* Execute the callback, storing the gradient and returning the value */
    vg = alloc(n*sizeof(double)/sizeof(void*), Double_array_tag);
    for (i = 0; i < n; i++) { Store_double_field(vg, i, 0.0); } 
    vret = callback2(*f, vx, vg);
    *pfx = Double_val(vret);
    /* Extract the gradient and return the value */
    for (i = 0; i < n; i++) { g[i] = Double_field(vg, i); }
    CAMLreturn0;
}

struct func_and_vec {
    value* f;
    double* x;
};

double val_and_grad2(void* instance, const double* x, double* g,
        int n, double step) {
    double fx;
    struct func_and_vec* fnx = (struct func_and_vec*)instance;
    val_and_grad(fnx->f, fnx->x, x, g, n, step, &fx);
    return fx;
}

CAMLprim value c_lbfgs(value f, value vx, value vc, value vxtol, value vmaxiter)
{
    CAMLparam5 (f, vx, vc, vxtol, vmaxiter);
    CAMLlocal1 (ret);
    int n = Wosize_val(vx);
    double c = Double_val(vc);
    double xtol = Double_val(vxtol);
    int maxiter = Int_val(vmaxiter);
    double* x = (double*)malloc(sizeof(double) * n);
    double fx;
    int i;
    int x_okay = 1;
    int errcode;
    struct func_and_vec fng;
    lbfgs_parameter_t param;

    for (i = 0; i < n; i++) { x[i] = Double_field(vx, i); }

    /* Initialize struct to hold data */
    fng.f = &f;
    fng.x = (double*)malloc(sizeof(double) * n);
    memcpy(fng.x, x, sizeof(double) * n);

    lbfgs_parameter_init(&param);
    param.xtol = xtol;
    param.max_iterations = maxiter;
    param.orthantwise_c = c;
    if (c > 0.0) {
        param.linesearch = LBFGS_LINESEARCH_BACKTRACKING;
    }
    errcode = lbfgs(n, x, &fx, val_and_grad2, NULL, (void*)&fng, &param); 
    for (i = 0; i < n; i++) {
        if (!finite(x[i])) {
          x_okay = 0;
        }
    }
    for (i = 0; i < n; i++) {
        Store_double_field(vx, i, x_okay ? x[i] : fng.x[i]);
    }
    free(x);
    free(fng.x);
    ret = caml_alloc(2, 0);
    Store_field(ret, 0, Val_int(errcode));
    Store_field(ret, 1, caml_copy_double(fx));
    fflush(stdout);
    CAMLreturn(ret);
}

